#%% init project env
import EsyPro
import json
import logging
import sqlite3
from datetime import datetime

sfs = EsyPro.ScriptResultManager('zqf', locals(), version=0)

#%%
import openai

# 从 JSON 文件加载 API 密钥
with open(sfs.path_of('config.json'), 'r') as config_file:
    config = json.load(config_file)
openai.api_key = config["openai_api_key"]

# 设置基础 URL
openai.base_url = "http://activity.scnet.cn:61080/v1/"

# 配置日志记录
logging.basicConfig(
    filename=sfs.path_of('chat_log.txt'),
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    encoding='utf-8'
)

class ChatDatabase:
    def __init__(self, db_path):
        """
        初始化数据库连接
        :param db_path: 数据库文件路径
        """
        self.db_path = db_path
        self._init_db()
        logging.info(f"初始化数据库连接: {db_path}")

    def _init_db(self):
        """初始化数据库表结构"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS conversations (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    model TEXT NOT NULL
                )
            ''')
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS messages (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    conversation_id INTEGER,
                    role TEXT NOT NULL,
                    content TEXT NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (conversation_id) REFERENCES conversations (id)
                )
            ''')
            conn.commit()
            logging.info("数据库表结构初始化完成")

    def create_conversation(self, model):
        """
        创建新的对话
        :param model: 使用的模型名称
        :return: 对话ID
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('INSERT INTO conversations (model) VALUES (?)', (model,))
            conversation_id = cursor.lastrowid
            logging.info(f"创建新对话: ID={conversation_id}, 模型={model}")
            return conversation_id

    def add_message(self, conversation_id, role, content):
        """
        添加消息到对话
        :param conversation_id: 对话ID
        :param role: 消息角色
        :param content: 消息内容
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute(
                'INSERT INTO messages (conversation_id, role, content) VALUES (?, ?, ?)',
                (conversation_id, role, content)
            )

    def get_conversation_history(self, conversation_id):
        """
        获取对话历史
        :param conversation_id: 对话ID
        :return: 消息列表
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute(
                'SELECT role, content FROM messages WHERE conversation_id = ? ORDER BY created_at',
                (conversation_id,)
            )
            messages = [{"role": role, "content": content} for role, content in cursor.fetchall()]
            logging.info(f"加载对话历史: ID={conversation_id}, 消息数量={len(messages)}")
            return messages

class MultiTurnChat:
    def __init__(self, model="QwQ-32B"):
        """
        初始化多轮对话类
        :param model: 使用的模型名称
        """
        self.model = model
        self.messages = []
        self.db = ChatDatabase(sfs.path_of('chat_history.db'))
        self.conversation_id = self.db.create_conversation(model)
        logging.info(f"初始化对话系统: 模型={model}")

    def add_user_message(self, content):
        """
        添加用户消息到对话历史
        :param content: 用户输入的内容
        """
        self.messages.append({"role": "user", "content": content})
        self.db.add_message(self.conversation_id, "user", content)

    def get_reply(self):
        """
        调用 OpenAI 接口获取 AI 回复
        :return: AI 回复的内容
        """
        try:
            completion = openai.chat.completions.create(
                model=self.model,
                messages=self.messages,
                stream=False,  # 设置为非流式输出
            )
            assistant_reply = completion.choices[0].message.content
            
            # 处理回复内容，去除内部思考过程
            if "</think>" in assistant_reply:
                assistant_reply = assistant_reply.split("</think>")[-1].strip()
            
            self.messages.append({"role": "assistant", "content": assistant_reply})
            self.db.add_message(self.conversation_id, "assistant", assistant_reply)
            return assistant_reply
        except Exception as e:
            error_msg = f"调用 OpenAI 接口失败: {e}"
            logging.error(error_msg)
            raise RuntimeError(error_msg)

    def load_conversation(self, conversation_id):
        """
        加载指定的对话历史
        :param conversation_id: 对话ID
        """
        self.conversation_id = conversation_id
        self.messages = self.db.get_conversation_history(conversation_id)
        logging.info(f"切换到对话: ID={conversation_id}")

#%% main
if __name__ == '__main__':
    print(f'start {__file__}')
    logging.info("启动对话系统")

    sfs.path_of().ensure()

    try:
        # 测试用例1：基本对话功能
        print("\n=== 测试用例1：基本对话功能 ===")
        chat1 = MultiTurnChat()
        print(f"创建新对话，ID: {chat1.conversation_id}")
        
        test_messages = [
            "你好，请介绍一下你自己",
            "你能做什么？",
            "给我讲个笑话"
        ]
        
        for msg in test_messages:
            print(f"\n用户: {msg}")
            chat1.add_user_message(msg)
            reply = chat1.get_reply()
            print(f"助手: {reply}")
        
        # 测试用例2：加载历史对话
        print("\n=== 测试用例2：加载历史对话 ===")
        chat2 = MultiTurnChat()
        print(f"创建新对话，ID: {chat2.conversation_id}")
        
        # 添加一些消息
        chat2.add_user_message("这是第二条对话")
        chat2.get_reply()
        
        # 加载第一条对话
        print(f"\n加载对话 {chat1.conversation_id}")
        chat2.load_conversation(chat1.conversation_id)
        
        # 继续对话
        chat2.add_user_message("我们继续之前的对话")
        reply = chat2.get_reply()
        print(f"\n用户: 我们继续之前的对话")
        print(f"助手: {reply}")

    except Exception as e:
        print(f"\n发生错误: {e}")
        logging.error(f"程序运行出错: {e}")
    finally:
        print("\n=== 测试结束 ===")
        print(f'end {__file__}')
        logging.info("对话系统关闭")
